#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   config.py
@Time    :   2024-08-04 19:38
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   定时调整虚拟机内存
"""

import logging
import os
import subprocess
import sys
import toml
from argparse import ArgumentParser
from datetime import datetime
from os import path
from sys import exit
from time import sleep

# 创建一个自定义的日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# 创建一个自定义的日志处理器，设置其输出格式
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d | %(funcName)s | %(message)s')
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(formatter)
logger.addHandler(handler)


def list_files(directory):
	"""
	列出指定目录下的所有文件
	"""
	file_list = []
	# 使用glob模块获取指定目录下的所有文件
	for root, dirs, files in os.walk(directory):
		for name in files:
			# print(os.path.join(root, name))
			file_list.append(os.path.join(root, name))
	return file_list


class QmSetMen:
	def __init__(self, config_list: str, debug: bool = False):
		"""

		:param config_list:
		:param debug:
		"""
		self.config_list = config_list
		self.debug = debug
		self.txt_config = {}
		self.txt_config_list = {}
		self.log_file = None
		# 当前时间:分
		self.current_min = 0
		# 当前时间:时
		self.current_hour = 0
		if self.debug:
			logger.debug("开始实例化")
		# 设置实例配置
		self.toml_name = ""
		self.toml_id = 220
		self.toml_mem_size = 0
		# 开启操作
		self.toml_start_hour = 0
		self.toml_on_delayed = 0

	def get_time(self):
		"""
		获取当前时间
		:return:
		"""
		if self.debug:
			logger.debug("正在获取当前时间")
		current = datetime.now()
		self.current_hour = current.hour
		self.current_min = current.minute
		if self.debug:
			logger.debug(f"当前小时: {self.current_hour} - 当前分钟: {self.current_min}")

	def read_list(self, file):
		"""
		读取配置文件
		:param file: 需要读取的主机配置文件
		:return:
		"""
		# 判断文件后缀
		if not file.endswith(".toml"):
			logger.warning("文件后缀错误,请检查文件后缀是否为.toml")
			return False
		if self.debug:
			logger.debug(f"正在读取配置文件: {file}")
		try:
			with open(file, "r", encoding='utf-8') as f:
				self.txt_config_list = toml.load(f)
				return True
		except FileNotFoundError as e:
			logger.error(f"配置文件不存在: {e}")
			return False
		except toml.TomlDecodeError as e:
			logger.error(f"读取配置文件失败,文件格式错误: {e}")
			return False
		except Exception as e:
			print(file)
			logger.error(f"读取配置文件失败: {e}")
			return False

	def result_log(self, result: str):
		"""
		输出结果
		:return:
		"""
		if self.debug:
			logger.debug("正在输出结果")
		if self.log_file:
			logger.info(f"日志文件: {self.log_file}")
		else:
			logger.info("日志文件: None")
			return
		if not path.isfile(self.log_file):
			os.system(f"touch {self.log_file}")
		try:
			with open(file=self.log_file, mode='a+', encoding='utf-8') as w:
				w.write(result)
		except Exception as e:
			logger.error(f"写入日志文件失败: {e}")

	def qm_status_to_stop(self, Id: int):
		"""
		将虚拟机关闭
		:param Id: 虚拟机ID
		:return:
		"""
		if self.qm_get_status(Id) != 0:
			if self.debug:
				logger.debug(f"正在关机: {Id}")
			# 开机
			ops = "关机"
			cmd_op = "stop"
			if self.debug:
				logger.debug(f"正在{ops}: {Id}")
			if os.system(f"qm {cmd_op} {Id}") == 0:
				logger.info(f"{ops}成功: {Id}")
				self.result_log(f"{datetime.now()} - {Id}  - {ops}成功\n")
				return True
			else:
				logger.error(f"{ops}失败: {Id}")
				self.result_log(f"{datetime.now()} - {Id}  - {ops}失败\n")
		else:
			logger.info(f"当前状态: {Id} - 不需要执行开机指令")
			return True
		return False

	def qm_status_to_start(self, Id: int):
		"""
		将虚拟机启动
		:param Id: 虚拟机ID
		:return:
		"""
		# 检查虚拟机状态是否为关闭，只有在关闭状态下才尝试启动
		if self.qm_get_status(Id) == 0:
			# 如果开启调试模式，则记录正在尝试开机的日志
			if self.debug:
				logger.debug(f"正在开机: {Id}")
			# 定义操作名称和命令
			ops = "开机"
			cmd_op = "start"
			# 如果开启调试模式，则记录正在执行的操作
			if self.debug:
				logger.debug(f"正在{ops}: {Id}")
			# 尝试执行开机命令，根据返回值判断操作是否成功
			if os.system(f"qm {cmd_op} {Id}") == 0:
				logger.info(f"{ops}成功: {Id}")
				self.result_log(f"{datetime.now()} - {Id}  - {ops}成功\n")
			else:
				logger.error(f"{ops}失败: {Id}")
				self.result_log(f"{datetime.now()} - {Id}  - {ops}失败\n")
		else:
			# 如果虚拟机不在关闭状态，则记录不需要执行开机指令的信息
			logger.info(f"当前状态: {Id} - 不需要执行开机指令")

	def qm_get_status(self, Id: int):
		"""
		获取虚拟机状态
		:param Id: 虚拟机ID
		:return: 1: 已启动/ 2: 已停止/ 3: 获取失败
		"""
		# 在调试模式下记录正在获取的虚拟机状态
		if self.debug:
			logger.debug(f"正在获取状态: {Id}")
		# 构建获取虚拟机状态的命令，使用awk提取状态信息
		cmd = str(f"qm status {Id} " + " | awk '{print $2}'")
		# 执行命令并获取结果
		res = subprocess.getstatusoutput(cmd)
		# 检查命令执行结果
		if res[0] == 0:
			# 判断虚拟机状态并返回相应代码
			if res[1].strip().lower() == "stopped":
				return 0
			else:
				return 1
		else:
			# 获取状态失败时记录错误日志
			logger.error(f"获取状态失败: {Id}")
			# 返回获取失败的错误码
			return 404

	def qm_set_cpu(self, Id, cpus: int):
		"""
		设置虚拟机CPU
		:param cpus: 设置CPU数量
		:param Id: 虚拟机ID
		:return:
		"""
		if self.debug:
			logger.debug(f"正在设置CPU数量: {Id}")
		# 获取当前CPU内核数量
		cmd = str("qm config %d | grep cores| awk '{print $2}'" % Id)
		if self.debug:
			print(cmd)
		res = subprocess.getstatusoutput(cmd)
		if res[0] == 0:
			res_out = res[1].strip().lower()
			try:
				size_ = int(res_out)
			except Exception as e:
				logger.error(f"获取CPU内核数量失败: {e}")
				return
			if self.debug:
				logger.debug(f"当前CPU内核数量: {res_out}")
				logger.debug(f"目标CPU内核数量: {cpus}")
			if int(size_) == int(cpus):
				logger.info(f"当前CPU内核数量: {res_out} - 目标CPU内核数量: {cpus} - 不需要修改")
				return
			if not self.qm_status_to_stop(Id):
				logger.error(f"停止虚拟机失败: {Id}")
				return
			cmd = str(f'qm set  {Id} --cores {cpus}')
			if self.debug:
				print(cmd)
				logger.debug(f"正在设置CPU数量: {Id}")
			if os.system(cmd) != 0:
				logger.error(f"设置CPU数量失败: {Id}")
				return
		else:
			logger.error(f"获取CPU内核数量失败: {Id}")

	def qm_set_mem(self, Id, Size: int):
		"""
		设置虚拟机内存大小。

		:param Id: 虚拟机ID。
		:param Size: 目标内存大小。
		"""
		# 初始化内存大小变量
		size_ = 0

		# 如果开启调试模式，则打印当前操作信息
		if self.debug:
			logger.debug(f"正在获取当前配置: {Id}")

		# 构建获取当前内存配置的命令
		cmd = str("qm config %d | grep memory| awk '{print $2}'" % Id)

		# 如果开启调试模式，则打印当前内存大小和命令信息
		if self.debug:
			print(size_)
			print(cmd)

		# 执行命令并获取结果
		res = subprocess.getstatusoutput(cmd)

		# 如果命令执行成功
		if res[0] == 0:
			# 移除结果两端的空白字符并转换为小写
			res_out = res[1].strip().lower()

			# 尝试将结果转换为整数
			try:
				size_ = int(res_out)
			except Exception as e:
				# 如果转换失败，则记录错误并返回
				logger.error(f"获取内存大小失败: {e}")
				return

			# 如果开启调试模式，则打印当前和目标内存大小
			if self.debug:
				logger.debug(f"当前大小: {res_out}")
				logger.debug(f"目标大小: {Size}")

			# 如果当前内存大小与目标内存大小相同，则记录信息并返回
			if int(size_) == int(Size):
				logger.info(f"当前内存大小: {size_} - 目标内存大小: {Size} - 不需要修改")
			else:
				# 如果当前内存大小与目标内存大小不同，则尝试停止虚拟机
				if not self.qm_status_to_stop(Id):
					# 如果停止虚拟机失败，则记录错误并返回
					logger.error(f"停止虚拟机失败: {Id}")
					return

				# 等待一段时间以确保虚拟机停止
				sleep(20)

				# 构建设置内存大小的命令
				cmd = str(f'qm set  {Id} --memory {Size}')

				# 如果开启调试模式，则打印当前操作信息
				if self.debug:
					print(cmd)
					logger.debug(f"正在设置内存大小: {Id}")

				# 执行设置内存大小的命令
				if os.system(cmd) != 0:
					# 如果设置内存大小失败，则记录错误并返回
					logger.error(f"设置内存大小失败: {Id}")
					return
		else:
			# 如果获取状态失败，则记录错误
			logger.error(f"获取状态失败: {Id}")
			print(res[1])

		# 等待一段时间以确保设置生效
		sleep(20)

	def check_time(self, start_hour, start_min, end_hour, end_min):
		"""
		判断时间是否在配置文件设置的时间范围内
		:param start_hour:
		:param start_min:
		:param end_hour:
		:param end_min:
		:return:
		"""
		# 格式化成时间对象
		start_time = datetime.strptime(f"{start_hour}:{start_min}", "%H:%M")
		end_time = datetime.strptime(f"{end_hour}:{end_min}", "%H:%M")
		# 获取当前时间对象，仅获取小时和分钟
		start_time = start_time.strftime("%H:%M")
		end_time = end_time.strftime("%H:%M")
		now_time = datetime.now().strftime("%H:%M")
		if self.debug:
			logger.debug(f"当前时间: {now_time}")
		# 如果当前时间在配置文件设置的时间范围内，则返回True，否则返回False
		if start_time <= now_time <= end_time:
			if self.debug:
				logger.debug(f"当前时间在配置文件设置的时间范围内: {start_time} - {end_time}")
			return True
		else:
			if self.debug:
				logger.debug(f"当前时间不在配置文件设置的时间范围内: {start_time} - {end_time}")
			return False

	def list_config(self):
		"""
		读取列表主机配置文件
		:return:
		"""
		# 判断文件夹是否存在且是否有文件
		if path.isdir(self.config_list):
			fl = list_files(self.config_list)
			if fl:
				if self.debug:
					# 打印文件数量
					logger.debug(f"当前配置文件数量: {len(fl)}")
				for file in fl:
					if not self.read_list(file=file):
						continue
					# 获取需要操作的主机ID列表
					data_list = self.txt_config_list["list"]
					host_list = data_list["list"]
					hour_start_ = data_list["start_hour"]
					min_start_ = data_list["start_minute"]
					hour_end_ = data_list["end_hour"]
					min_end = data_list["end_minute"]
					size_ = data_list["size"]
					cpus_ = data_list["cpus"]
					enable_ = data_list["enabled"]
					# 检查时间是否在设定的范围内，且配置项已启用
					if self.check_time(hour_start_, min_start_, hour_end_, min_end) and enable_:
						stop_cmd = "systemctl stop com.liumou.pve.crontab.service "
						os.system(stop_cmd)
						for i in host_list:
							if self.debug:
								logger.debug(f"正在检测并设置虚拟机配置: {i}")
							self.qm_set_mem(i, size_)
							self.qm_set_cpu(i, cpus_)
							self.qm_status_to_start(i)
						start_cmd = "systemctl start com.liumou.pve.crontab.service "
						os.system(start_cmd)

			else:
				logger.warning(f"配置文件目录为空: {self.config_list}")
		else:
			if self.debug:
				logger.debug(f"配置文件目录不存在: {self.config_list}")

	def start(self):
		while True:
			self.list_config()


def install():
	"""
	安装/更新配置文件和相关资源。

	该函数负责从指定的URL地址下载配置文件和资源文件，并将它们保存到系统中的特定目录。
	如果文件已存在，将先删除再下载。下载成功后会对相关文件进行权限设置以确保其可执行性。
	安装或更新完成后，会调用service函数并退出程序。

	:return: 无
	"""
	# 定义本地二进制文件路径和URL地址，用于后续的文件操作和下载
	memory_bin = path.join("/usr/bin/com.liumou.pve.qm.config")
	memory_url = "https://gitee.com/liumou_site/pvet/raw/master/model/qm/config.py"
	# 记录安装/更新开始的信息
	logger.info("正在进行安装/更新...")
	# 定义内存配置文件的URL，用于下载最新的配置
	url_list = 'https://gitee.com/liumou_site/pvet/raw/master/model/qm/memory.toml'
	# 创建存储配置文件的目录，确保路径存在
	os.system("mkdir -p /opt/apps/com.liumou.pve.qm.config/")
	# 定义示例列表的本地路径和URL，用于下载示例数据
	demo_list = "/opt/apps/com.liumou.pve.qm.config/demo_list.toml"
	# 创建一个字典，映射本地文件路径到对应的下载URL
	url_dict = {demo_list: url_list, memory_bin: memory_url}
	# 遍历字典中的每个文件，检查其是否存在，如果存在则删除，然后尝试下载
	for file in url_dict:
		if path.exists(file):
			logger.info("模板配置文件已存在")
			os.system(f"rm -f {file}")
		# 使用wget命令下载文件，检查下载结果
		res = subprocess.getstatusoutput(f'wget -O {file} {url_dict[file]}')
		if res[0] == 0:
			logger.info(f"[ {file} ]配置文件下载成功")
		else:
			# 如果下载失败，记录错误信息并退出程序
			logger.info(f"[ {file} ]配置文件下载失败")
			print(res[1])
			exit(1)
	# 下载成功后，设置二进制文件的权限，使其可执行
	os.system(f"chmod 0755 {memory_bin}")
	# 记录安装/更新完成的信息
	logger.info("安装/更新成功")
	# 调用service函数进行后续操作，并退出程序
	service()
	exit(0)


def service():
	"""
	配置和启动一个Systemd服务，该服务使用Proxmox VE（PVE）的定时工具。

	该函数执行以下操作：
	1. 检查是否存在的执行文件: /usr/bin/com.liumou.pve.qm.config。
	2. 若存在，创建服务配置文件: /etc/systemd/system/com.liumou.pve.qm.config.service。
	3. 重新加载Systemd配置，启用并重启创建的服务。

	:return: 无
	"""
	# 检查PVE定时工具的执行文件是否存在，不存在则输出错误并退出
	if not path.isfile("/usr/bin/com.liumou.pve.qm.config"):
		logger.error(f"执行文件不存在: /usr/bin/com.liumou.pve.qm.config")
		exit(1)

	# 定义服务配置文件的内容
	txt = """[Unit]
Description=This is a PVE timing tool
After=network.target

[Service]
Type=simple
User=root
ExecStart=/usr/bin/com.liumou.pve.qm.config
StandardOutput=journal
KillMode=control-group
Restart=on-failure
RestartPreventExitStatus=255
StandardError=inherit
SyslogIdentifier=com.liumou.pve.qm.config

[Install]
WantedBy=multi-user.target
Alias=com.liumou.pve.qm.config.service
"""

	# 尝试写入服务配置文件，失败则输出错误并退出
	try:
		with open(file="/etc/systemd/system/com.liumou.pve.qm.config.service", mode='w', encoding='utf-8') as w:
			w.write(txt)
	except Exception as e:
		logger.error(f"写入服务文件失败: {e}")
		exit(1)

	# 重新加载Systemd配置，若成功则启用并重启服务，检查启动结果并相应输出日志、退出
	if os.system("systemctl daemon-reload") == 0:
		os.system("systemctl enable com.liumou.pve.qm.config.service")
		if os.system("systemctl restart com.liumou.pve.qm.config.service") == 0:
			logger.info("服务启动成功")
			exit(0)
		else:
			logger.error("服务启动失败...")
			exit(3)
	else:
		logger.error("服务启动失败")
		exit(2)


# 当作为主程序运行时，初始化默认配置列表路径和日志目录
if __name__ == '__main__':
	conf_default_list = "/opt/apps/com.liumou.pve.qm.config/config_list/"
	os.system(f"mkdir -p /opt/apps/com.liumou.pve.qm.config/log/ {conf_default_list}")
	# 创建命令行参数解析器，用于解释脚本运行时的参数
	arg = ArgumentParser(description='当前脚本版本:1.0', prog="PVE虚拟机/容器内存定时调整")
	h = f"设置列表配置文件路径,默认: {conf_default_list}"
	# 添加命令行参数：配置列表路径
	arg.add_argument('-cl', '--conf_list', type=str, help=h, default=conf_default_list, required=False)
	# 添加命令行参数：强制更新配置文件选项
	arg.add_argument('-f', action='store_true', help='强制更新配置文件')
	# 添加命令行参数：调试模式选项
	arg.add_argument('-debug', action='store_true', help='开启调试模式')
	# 添加命令行参数：安装/更新服务选项
	arg.add_argument('-install', action='store_true', help='安装/更新服务')
	# 解析命令行参数
	args = arg.parse_args()
	# 提取解析后的配置列表路径
	args_conf_list = args.conf_list
	# 提取解析后的安装/更新选项
	args_install = args.install
	# 提取解析后的调试模式选项
	args_debug = args.debug
	# 提取解析后的强制更新配置文件选项
	args_f = args.f
	# 如果开启了调试模式，记录相关信息
	if args_debug:
		logger.info("开启了调试模式")
	# 如果开启了安装/更新功能，执行安装/更新操作
	if args_install:
		logger.info("开启了安装/更新功能")
		install()
	# 根据解析后的参数，初始化并启动内存调整任务
	c = QmSetMen(debug=args_debug, config_list=args_conf_list)
	c.start()
